# File Name:    inference.py
# Version:      ver1_0
# Created:      2024/06/17
# Description:  本文件定义了用于在模型应用端进行推理，返回模型输出的流程
#               ★★★请在空白处填写适当的语句，将模型推理应用流程补充完整★★★
# -----------------------------------------------------------------------

import torch
from PIL import Image
from torchvision import transforms
import matplotlib.pyplot as plt
from model import CustomNet  # 导入模型类


def inference(image_path, model, device):
    """定义模型推理应用的流程。
    :param image_path: 输入图片的路径
    :param model: 训练好的模型
    :param device: 模型推理使用的设备，即使用哪一块CPU、GPU进行模型推理
    """
    # 将模型置为评估（测试）模式
    model.eval()

    # START----------------------------------------------------------
    # 定义图像预处理转换
    transform = transforms.Compose([
        transforms.Resize((224, 224)),  # 根据模型输入大小调整
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])

    # 打开并预处理图像
    image = Image.open(image_path).convert('RGB')
    input_tensor = transform(image).unsqueeze(0).to(device)  # 添加批次维度并移至设备

    # 执行模型推理
    with torch.no_grad():
        output = model(input_tensor)
        # 应用softmax获取类别概率
        probabilities = torch.nn.functional.softmax(output, dim=1)
        # 获取预测的类别和置信度
        confidence, predicted_class = torch.max(probabilities, 1)

    # 打印结果
    print(f"预测类别: {predicted_class.item()}")
    print(f"置信度: {confidence.item() * 100:.2f}%")

    # 可视化预测结果
    plt.figure(figsize=(10, 4))

    # 显示图像
    plt.subplot(1, 2, 1)
    plt.imshow(image)
    plt.title(f"预测结果: {predicted_class.item()}")
    plt.axis('off')

    # 显示概率分布
    plt.subplot(1, 2, 2)
    classes = [str(i) for i in range(probabilities.size(1))]
    plt.bar(classes, probabilities.cpu().numpy()[0])
    plt.xlabel('Class')
    plt.ylabel('Probability')
    plt.title('Categorical Distribution')
    plt.ylim(0, 1.0)

    plt.tight_layout()
    plt.show()

    # END------------------------------------------------------------


if __name__ == "__main__":
    # 指定图片路径
    image_path = "./images/test/signs/img_0006.png"

    # 创建模型实例
    model = CustomNet(num_classes=6)  # 根据你的模型修改类别数

    # 设置设备（在加载模型之前定义device）
    if torch.cuda.is_available():
        device = torch.device("cuda")
        print("使用GPU进行推理")
    else:
        device = torch.device("cpu")
        print("使用CPU进行推理")

    # 加载模型权重
    try:
        model_state_dict = torch.load('./models/best_model.pth', map_location=device)
        model.load_state_dict(model_state_dict)
        print("模型权重加载成功")
    except Exception as e:
        print(f"模型加载失败: {e}")
        exit(1)

    # 将模型移至设备
    model.to(device)

    # 显示图片，输出预测结果
    inference(image_path, model, device)